import os
import subprocess
from threading import Thread
import time

import cv2
import pyautogui
from pynput.keyboard import Listener
from pynput.keyboard import Key
from ultralytics import YOLO

wd = os.getcwd()


class VideoRecorder(Thread):
	popen: subprocess.Popen
	mp4Path = ""

	def __init__(self, name=None):
		Thread.__init__(self, name=name)

	def run(self):
		self.mp4Path = os.path.join(wd, "test.mp4")
		sw, sh = pyautogui.size()
		screenResolution = "{}x{}".format(sw, sh)
		cmd = "{} -f gdigrab -i desktop -y -r 30 -video_size {} {}" \
			.format(os.path.join(wd, "ffmpeg/bin/ffmpeg.exe"), screenResolution, self.mp4Path)
		self.popen = subprocess.Popen(cmd, shell=True, stdin=subprocess.PIPE)

		with Listener(on_press=self.stop_record) as listener:
			listener.join()

	def stop_record(self, key):
		try:
			if key == Key.f7:
				self.popen.stdin.write("q".encode("GBK"))
				self.popen.communicate()
				self.popen.kill()
		except AttributeError:
			return


def getModel():
	model = YOLO(os.path.join(wd, "model.yaml"))
	model = YOLO(os.path.join(wd, "best.pt"))
	return model


def train(model):
	model.train(data=os.path.join(wd, "data.yaml"), epochs=1, patience=200, batch=-1, imgsz=640, device=0, workers=2)
	metrics = model.val()
	# success = model.export(format="onnx", opset=12, simplify=True, dynamic=False, imgsz=640)
	return model


if __name__ == '__main__':
	try:
		model = getModel()

		recorder = VideoRecorder("video")
		recorder.start()

		while recorder.mp4Path == "":
			time.sleep(0.025)

		while recorder.mp4Path != "":
			cap = cv2.VideoCapture(recorder.mp4Path)
			while True:
				ok, frame = cap.read()
				if not ok:
					break
				tp = os.path.join(wd, "tmp/") + str(int(time.time() * 1000000)) + ".jpg"
				cv2.imwrite(tp, frame)

				results = model(source=tp, save=False, name=os.path.join(wd, "runs/detect"))
				os.remove(tp)

				if len(results) == 0:
					continue
				boxes = results[len(results) - 1:][0]
				names = boxes.names
				boxes = boxes.boxes

				r = ""
				for i, cls in enumerate(boxes.cls):
					label = names[int(cls)]
					confidence = round(float(boxes.conf[i]), 2)
					x = int(boxes.xywh[i][0])
					y = int(boxes.xyxy[i][3])

					r += f"{label}:{confidence}:{x},{y}\n"

				predictFile = open(os.path.join(wd, "runs/detect/predict"), "w")
				predictFile.write(r)
				predictFile.flush()
				predictFile.close()

			cap.release()
			cv2.destroyAllWindows()

	except Exception as e:
		print(f"{e}")

# cxfreeze -c train_t.py --target-dir dist
